

from sketchz import *

prog_render2d2 = {
    "vertex_shader": '''
#version 330 core

layout (location = 0) in int in_id;

vec3[4] pos_list = vec3[](
    vec3(-1.0, 1.0, 0),
    vec3(1.0, 1.0, 0),
    vec3(-1.0, -1.0, 0),
    vec3(1.0, -1.0, 0)
);

vec2[4] uv_list = vec2[](
    vec2(0, 1.0),
    vec2(1.0, 1.0),
    vec2(0, 0),
    vec2(1.0, 0)
);

layout (std140) uniform Common {
    float gray;
    float alpha;
    float z;
    float font_gray;
    float win_width;
    float win_height;
    float n1;
    float n2;
};

out vec2 uv;

void main() {
    gl_Position = vec4(pos_list[in_id], 1.0);
    gl_Position.z = z;
    uv = vec2(uv_list[in_id].x, 1 - uv_list[in_id].y);
}
 ''',
    "fragment_shader": '''
#version 330 core
layout (location = 0) out vec4 fragColor;

layout (std140) uniform Common {
    float gray;
    float alpha;
    float z;
    float font_gray;
    float win_width;
    float win_height;
    float n1;
    float n2;
};


uniform sampler2D texture0;
uniform sampler2D texture1;

in vec2 uv;

vec4 alphaBlend(vec4 top, vec4 bottom)
{
    vec3 color = (top.rgb * top.a) + (bottom.rgb * (1 - top.a));
    float alpha = top.a + bottom.a * (1 - top.a);

    return vec4(color, alpha);
}

void main()
{
    vec4 rgba = texture(texture0, uv);



    rgba.a = rgba.a * font_gray;
    rgba = alphaBlend(rgba, vec4(gray, gray, gray, alpha));

    vec2 uv2 = uv * vec2(win_width, win_height) - vec2(300, 300);
    vec2 uv3 = uv2 / vec2(100, 100)/ vec2(15, 33);
    float m = texture(texture1, uv3).r;
    //rgba.rgb += vec3(m, m, m);
    vec4 temp = alphaBlend(rgba, vec4(1, 1, 1, 1));
    temp.rgb = vec3(1, 1, 1) - temp.rgb;

    rgba = mix(rgba, temp, m);
    fragColor = vec4(rgba.rgb, rgba.a);
    //fragColor = vec4(rgba.a, 0, 0, 1);
}

            '''
}

class Render2d2(Drawable):
    def __init__(self):
        self.pipeline  = None
        self.uniform_buffer = None
        self.texture = None
        self.texture2 = None

    def u(self, window, sdlsurf=False):
        assert hasattr(window, "render")
        assert hasattr(window, "update")
        assert hasattr(window, "mainloop")

        try:
            self.uniform_buffer = window.ctx.buffer(size=32, dynamic=True)
        except:
            self.uniform_buffer = window.ctx.buffer(size=32, uniform=True)
        self.texture = window.ctx.image(window.size, 'rgba8unorm')
        self.texture2 = window.ctx.image((100, 100), 'rgba8unorm')

        # ~ arr = create_screen_quad()
        # ~ buf = window.ctx.buffer(arr.tobytes())
        buf = window.ctx.buffer(np.array([0, 2, 3, 0, 3, 1]))
        self.pipeline = window.ctx.pipeline(
            vertex_shader = prog_render2d2["vertex_shader"],
            fragment_shader = prog_render2d2["fragment_shader"],
            framebuffer = [window.img, window.imgs],
            layout=[
                {'name': 'Common', 'binding': 0},
                {'name': 'texture0', 'binding': 0},
                {'name': 'texture1', 'binding': 1},
            ],
            resources=[
                {
                    'type': 'uniform_buffer',
                    'binding': 0,
                    'buffer': self.uniform_buffer,
                },
                {
                    'type': 'sampler',
                    'binding': 0,
                    'image': self.texture,
                },
                {
                    'type': 'sampler',
                    'binding': 1,
                    'image': self.texture2,
                    'min_filter': 'nearest',
                    'mag_filter': 'nearest',
                },
            ],
            vertex_buffers = zengl.bind(buf, "1i", 0,),
            vertex_count = buf.size // zengl.calcsize('1i'),
            blend = {
                'enable': True,
                'src_color': 'src_alpha',
                'dst_color': 'one_minus_src_alpha',
                'src_alpha': 'one',
                'dst_alpha': 'one_minus_src_alpha',
            },
        )

        # ~ v = Quaternion().matrix44 *  Matrix44.from_translation(-U3(0, -10, 0))
        # ~ p = Matrix44.orthogonal_projection(-1, 1, -1, 1, 0, window.size[0])
        # ~ am = (p * v).astype("f4")

        self.uniform_buffer.write(np.array([0, 0.1, -0.99, 1, window.size[0], window.size[1], 0, 0], dtype="f4").tobytes())
        # ~ self.uniform_buffer.write(struct.pack('16s', am.tobytes()))

    def draw(self):
        self.pipeline.render()



class Window2:
    def __init__(self, size=(800, 600)):
        sx, sy = self.size = size
        self.size_half = (size[0] // 2, size[1] // 2)
        self.camera = Camera(U3(0, 0, 0), xy=sx/sy)   # 默认摄像机
        self.draw3d_list = []

        self.key_down_callback = []
        self.key_up_callback = []
        self.mouse_down_callback = []
        self.mouse_up_callback = []
        self.mouse_move_callback = []
        self.wheel_callback = []
        self.update_callback = []
        self.text_callback = []


        # ~ sdl2.SDL_SetHint(sdl2.SDL_HINT_OPENGL_ES_DRIVER, b"0")
        sdl2.SDL_GL_SetAttribute(sdl2.SDL_GL_CONTEXT_MAJOR_VERSION, 4)
        sdl2.SDL_GL_SetAttribute(sdl2.SDL_GL_CONTEXT_MINOR_VERSION, 3)
        sdl2.SDL_GL_SetAttribute(sdl2.SDL_GL_CONTEXT_PROFILE_MASK, sdl2.SDL_GL_CONTEXT_PROFILE_CORE)

        sdl2.SDL_GL_SetAttribute(sdl2.SDL_GL_DOUBLEBUFFER, 1)
        sdl2.SDL_GL_SetAttribute(sdl2.SDL_GL_DEPTH_SIZE, 24)
        err = sdl2.SDL_Init(sdl2.SDL_INIT_VIDEO)
        if err != 0:
            print(sdl2.SDL_GetError())

        self.window = window = sdl2.SDL_CreateWindow(b"sketchz",
            sdl2.SDL_WINDOWPOS_UNDEFINED,
            sdl2.SDL_WINDOWPOS_UNDEFINED, sx, sy,
            sdl2.SDL_WINDOW_OPENGL)

        if not window:
            print(sdl2.SDL_GetError())
            exit()

        self._context = sdl2.SDL_GL_CreateContext(window)
        self.ctx = zengl.context()

        self.img = self.ctx.image((sx, sy), "rgba32float")
        self.img.clear_value = (1, 1, 1, 1)
        self.imgh = self.ctx.image((sx, sy), "rgba32float")     # 很好，同时输出法线和未正规化的深度
        self.imgs = self.ctx.image((sx, sy), "depth32float")


        self.screen = sdl2.ext.surface._create_surface(self.size).contents

        self.render2d = Render2d2()

        self.render2d.u(self)





    def o(self, p):
        x, y = p
        a, b = self.size_half
        return V2(x * 100 + a, y * -100 + b)


    def o_inv(self, p):
        x, y = p
        a, b = self.size_half
        return V2((x - a) / 100, (b - y) / 100)


    def os(self, x):
        return x * 100


    def os_inv(self, x):
        return x / 100


    def __rrshift__(self, obj):
        if isinstance(obj, Drawable):
            self.draw3d_list.append(obj)
        elif isinstance(obj, Control):
            if hasattr(obj, "draw"):
                self.draw2d_list.append(obj)
            if hasattr(obj, "reg_event"):     # 并不能确定要注册多少种事件，并且事件可以临时删除
                obj.reg_event(self)
        elif isinstance(obj, Camera):
            self.camera = obj
        else:
            raise Exception(f" >> window fail, {obj}")

    def dispatch(self, e):
        if e.type == sdl2.SDL_MOUSEBUTTONDOWN:
            for callback in self.mouse_down_callback:
                r = callback(self, e.button)
                if r == True:
                    break
        elif e.type == sdl2.SDL_MOUSEBUTTONUP:
            for callback in self.mouse_up_callback:
                r = callback(self, e.button)
                if r == True:
                    break
        elif e.type == sdl2.SDL_MOUSEMOTION:
            for callback in self.mouse_move_callback:
                callback(self, e.motion)
        elif e.type == sdl2.SDL_KEYDOWN:
            for callback in self.key_down_callback:
                r = callback(self, e.key)
                if r == True:
                    break
        elif e.type == sdl2.SDL_KEYUP:
            for callback in self.key_up_callback:
                r = callback(self, e.key)
                if r == True:
                    break
        elif e.type == sdl2.SDL_MOUSEWHEEL :
            for callback in self.wheel_callback:
                r = callback(self, e.wheel)
                if r == True:
                    break
        elif e.type == sdl2.SDL_TEXTINPUT:
            for callback in self.text_callback:
                r = callback(self, e.text)
                if r == True:
                    break


    def render(self):
        self.ctx.new_frame()

        self.img.clear()
        self.imgs.clear()
        self.imgh.clear()

        for obj in self.draw3d_list:
            obj.draw(self.camera)       # 暂时不考虑灯光



        self.render2d.draw()

        self.img.blit()

        self.ctx.end_frame()

    def flush_screen(self):
        to_image(self.screen, self.render2d.texture)

    def update(self, dt):
        for callback in self.update_callback:
            callback(self, dt)

    def mainloop(self):
        running = True
        e = sdl2.SDL_Event()
        while running:
            while sdl2.SDL_PollEvent(ctypes.byref(e)) != 0:
                if e.type == sdl2.SDL_QUIT:
                    running = False
                self.dispatch(e)

            self.update(0.02)
            self.render()

            sdl2.SDL_GL_SwapWindow(self.window)
            sdl2.SDL_Delay(20)


        sdl2.SDL_DestroyWindow(self.window)
        sdl2.SDL_Quit()

